import time
import json
import uuid
import asyncio
from flask import Flask, request, jsonify

# ========= 任务管理 =========
tasks = {}

async def test_async(payload=None, url='', task_id=None):
    headers = {
        'Authorization': 'Bearer sk-VVJU25qYb5Bm3NHiIjYGevGfvQBZIj838sNC1D73FM0rqFpT',
        'Content-Type': 'application/json'
    }
    import aiohttp
    timeout_cfg = aiohttp.ClientTimeout(total=600)

    async with aiohttp.ClientSession(timeout=timeout_cfg) as session:
        async with session.post(url, headers=headers, json=payload) as response:
            response.raise_for_status()
            if response.status == 200:
                result = await response.json()
                response_content = result['choices'][0]['message']['content']
                tasks[task_id] = {
                    "status": "completed",
                    "response_content": response_content
                }
                return response_content
            else:
                tasks[task_id] = {
                    "status": "failed",
                    "response_content": None
                }
                return None

def start_background_loop(loop):
    asyncio.set_event_loop(loop)
    loop.run_forever()

# ========= 异步事件循环/后台线程支持 =========
import threading
background_loop = asyncio.new_event_loop()
t = threading.Thread(target=start_background_loop, args=(background_loop,), daemon=True)
t.start()

# ========= Flask 服务 =========
app = Flask(__name__)

@app.route('/chat', methods=['POST'])
def chat():
    """
    接收chat消息，异步处理，立即返回task_id
    """
    try:
        data = request.get_json(force=True)
        messages = data.get('messages')
        if not messages:
            return jsonify({"success": False, "msg": "Missing messages"}), 400

        url = "https://ismaque.org/v1/chat/completions"
        payload = {
            "model": "gemini-2.5-flash-image-vip",
            "stream": False,
            "messages": messages
        }

        task_id = str(uuid.uuid4())
        # 标记为运行中
        global tasks
        tasks[task_id] = {
            "status": "running",
            "response_content": None
        }
        # 启动异步后台任务
        # asyncio.create_task(test_async(payload=payload, url=url, task_id=task_id))
        asyncio.run_coroutine_threadsafe(test_async(payload=payload, url=url, task_id=task_id), background_loop)
        global tasks
        print(tasks)
        return jsonify({"success": True, "task_id": task_id})
    except Exception as e:
        return jsonify({"success": False, "msg": str(e)}), 500

@app.route('/query_chat', methods=['POST'])
def query_chat():
    global tasks
    print(tasks)
    """
    提交task_id查询结果
    """
    try:
        data = request.get_json(force=True)
        task_id = data.get("task_id")
        if not task_id:
            return jsonify({"success": False, "msg": "Missing task_id"}), 400
        task_info = tasks.get(task_id, {})
        status = task_info.get("status", "")
        if status == "completed":
            return jsonify({"success": True, "status": "completed", "result": task_info.get("response_content")})
        elif status == "running":
            return jsonify({"success": True, "status": "running", "result": None})
        elif status == "failed":
            return jsonify({"success": False, "status": "failed", "result": None})
        else:
            return jsonify({"success": False, "status": "not_found", "result": None})
    except Exception as e:
        return jsonify({"success": False, "msg": str(e)}), 500

if __name__ == '__main__':
    app.run("0.0.0.0", 20026, debug=False)
